package baidumodel.service;

import baidumodel.constant.ApiConstant;
import baidumodel.entity.chat.*;
import baidumodel.exception.BaiduModelException;
import baidumodel.util.HttpUtil;
import baidumodel.util.TokenUtil;
import cn.hutool.http.ContentType;
import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.NonNull;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;

import java.net.Proxy;
import java.util.Objects;
import java.util.concurrent.TimeUnit;


/**
 * 发送接口请求
 *
 * @Author wuchenxi
 * @Date 2023-07-23 17:46:49
 */
public class BaiduService {

    private OkHttpClient okHttpClient;
    private final String apiKey;
    private final String secretKey;

    public BaiduService(@NonNull String apiKey, @NonNull String secretKey) {
        this.apiKey = apiKey;
        this.secretKey = secretKey;
        this.okHttpClient(30, 30, 30, null);
    }

    public BaiduService(@NonNull String apiKey, @NonNull String secretKey, long connectTimeout, long writeTimeout, long readTimeout, Proxy proxy) {
        this.apiKey = apiKey;
        this.secretKey = secretKey;
        this.okHttpClient(connectTimeout, writeTimeout, readTimeout, proxy);
    }


    private void okHttpClient(long connectTimeout, long writeTimeout, long readTimeout, Proxy proxy) {
        OkHttpClient.Builder client = new OkHttpClient.Builder();
        client.connectTimeout(connectTimeout, TimeUnit.SECONDS);
        client.writeTimeout(writeTimeout, TimeUnit.SECONDS);
        client.readTimeout(readTimeout, TimeUnit.SECONDS);
        if (Objects.nonNull(proxy)) {
            client.proxy(proxy);
        }
        this.okHttpClient = client.build();
    }

    /**
     * 发起 ernieBot模型会话（同步接口）
     *
     * @param param
     * @param token 鉴权token
     * @return
     */
    public ErnieBotResponse ernieBot(ErnieBotParam param, String token) {
        return (ErnieBotResponse) this.chat(ApiConstant.ERNIE_BOT, param, token, ErnieBotResponse.class);
    }

    /**
     * 发起 ernieBot模型会话（异步接口）
     *
     * @param param
     * @param token 鉴权token
     * @return
     */
    public void ernieBotStream(ErnieBotParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(ApiConstant.ERNIE_BOT, param, token, eventSourceListener);
    }

    /**
     * 发起 ernieBotTurbo 模型会话（同步接口）
     *
     * @param param
     * @param token 鉴权token
     * @return
     */
    public ErnieBotTurboResponse ernieBotTurbo(ErnieBotTurboParam param, String token) {
        return (ErnieBotTurboResponse) this.chat(ApiConstant.ERNIE_BOT_TURBO, param, token, ErnieBotTurboResponse.class);
    }



    /**
     * 发起 ernieBotTurbo 模型会话（流式）
     *
     * @param param
     * @param token               鉴权token
     * @param eventSourceListener {@link baidumodel.listener.BaiduEventSourceListener}
     * @return
     */
    public void ernieBotTurboStream(ErnieBotTurboParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(ApiConstant.ERNIE_BOT_TURBO, param, token, eventSourceListener);
    }

    /**
     * 发起 ERNIE_BOT_4 模型对话（同步接口）
     *
     * @param param
     * @param token
     * @return
     */
    public ErnieBot4Response ernieBot4(ErnieBot4Param param, String token) {
        return (ErnieBot4Response) this.chat(ApiConstant.ERNIE_BOT_4, param, token, ErnieBot4Response.class);
    }

    /**
     * 发起 ERNIE_BOT_4 模型对话（流式）
     *
     * @param param
     * @param token               鉴权token
     * @param eventSourceListener {@link baidumodel.listener.BaiduEventSourceListener}
     * @return
     */
    public void ernieBot4Stream(ErnieBot4Param param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(ApiConstant.ERNIE_BOT_4, param, token, eventSourceListener);
    }

    /**
     * 发起 ernieBot模型会话（同步接口）
     *
     * @param param
     * @param token 鉴权token
     * @return
     */
    public Bloomz7bResponse bloomz7b(Bloomz7bParam param, String token) {
        return (Bloomz7bResponse) this.chat(ApiConstant.BLOOMZ_7B, param, token, Bloomz7bResponse.class);
    }

    /**
     * 发起 ernieBot模型会话（异步接口）
     *
     * @param param
     * @param token 鉴权token
     * @return
     */
    public void bloomz7bStream(Bloomz7bParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(ApiConstant.BLOOMZ_7B, param, token, eventSourceListener);
    }

    /**
     * 获取prompt模板信息
     *
     * @param param
     * @param token
     * @return
     */
    public PromptTemplateResponse promptTemplate(PromptTemplateParam param, String token) {
        if (param == null) {
            throw new BaiduModelException("参数异常：param不能为空");
        }
        StringBuilder url = new StringBuilder(ApiConstant.PROMPT_TEMPLATE).append(token)
                .append("&id=").append(param.getId());
        if (param.getValues() != null && !param.getValues().isEmpty()) {
            param.getValues().forEach((k, v) -> url.append("&").append(k).append("=").append(v));
        }
        String s = HttpUtil.get(url.toString());
        return JSONUtil.toBean(s, PromptTemplateResponse.class);
    }

    /**
     * 获取请求token
     *
     * @return token
     */
    public String getToken() {
        return TokenUtil.getToken(apiKey, secretKey);
    }

    /**
     * embeddingV1接口
     *
     * @param param
     * @param token
     * @return
     */
    public EmbeddingV1Response embeddingV1(EmbeddingV1Param param, String token) {
        String post = HttpUtil.post(ApiConstant.EMBEDDING_V1 + token, JSONUtil.toJsonStr(param));
        return JSONUtil.toBean(post, EmbeddingV1Response.class);
    }

    /**
     * Llama-2-70b-chat
     *
     * @param param
     * @param token
     * @return
     */
    public LlamaChatResponse Llama2_70bChat(Llama2ChatParam param, String token) {
        return (LlamaChatResponse) this.chat(String.format(ApiConstant.LLAMA_2_70B_CHAT), param, token, LlamaChatResponse.class);
    }

    /**
     * Llama-2-7b-chat
     *
     * @param param
     * @param token
     * @return
     */
    public LlamaChatResponse Llama2_7bChat(Llama2ChatParam param, String token) {
        return (LlamaChatResponse) this.chat(String.format(ApiConstant.LLAMA_2_7B_CHAT), param, token, LlamaChatResponse.class);
    }

    /**
     * Llama-2-13b-chat
     *
     * @param param
     * @param token
     * @return
     */
    public LlamaChatResponse Llama2_13bChat(Llama2ChatParam param, String token) {
        return (LlamaChatResponse) this.chat(String.format(ApiConstant.LLAMA_2_13B_CHAT), param, token, LlamaChatResponse.class);
    }

    /**
     * Llama-2-70b-chat（流式）
     *
     * @param param
     * @param token
     * @param eventSourceListener
     * @return
     */
    public void Llama2_70bChatStream(Llama2ChatParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(String.format(ApiConstant.LLAMA_2_70B_CHAT), param, token, eventSourceListener);
    }

    /**
     * Llama-2-7b-chat（流式）
     *
     * @param param
     * @param token
     * @param eventSourceListener
     * @return
     */
    public void Llama2_7bChatStream(Llama2ChatParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(String.format(ApiConstant.LLAMA_2_7B_CHAT), param, token, eventSourceListener);
    }

    /**
     * Llama-2-13b-chat（流式）
     *
     * @param param
     * @param token
     * @param eventSourceListener
     * @return
     */
    public void Llama2_13bChatStream(Llama2ChatParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(String.format(ApiConstant.LLAMA_2_13B_CHAT), param, token, eventSourceListener);
    }

    /**
     * Linly-Chinese-LLaMA-2-7B
     * Linly-Chinese-LLaMA-2-13B
     *
     * @param param
     * @param token
     * @return
     */
    public LinlyChineseLLaMAResponse linlyChineseLLaMA(LinlyChineseLLaMAParam param, String token) {
        return (LinlyChineseLLaMAResponse) this.chat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, LinlyChineseLLaMAResponse.class);
    }

    /**
     * ChatGLM2-6B
     *
     * @param param
     * @param token
     * @return
     */
    public ChatGLM26BResponse chatGLM26B(ChatGLM26BParam param, String token) {
        return (ChatGLM26BResponse) this.chat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, ChatGLM26BResponse.class);
    }

    /**
     * ChatGLM2-6B
     *
     * @param param
     * @param token
     * @param eventSourceListener
     * @return
     */
    public void chatGLM26BStream(ChatGLM26BParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, eventSourceListener);
    }

    /**
     * OpenLLaMA-7B
     *
     * @param param
     * @param token
     * @return
     */
    public OpenLLaMA7BResponse openLLaMA7B(OpenLLaMA7BParam param, String token) {
        return (OpenLLaMA7BResponse) this.chat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, OpenLLaMA7BResponse.class);
    }

    /**
     * OpenLLaMA-7B
     *
     * @param param
     * @param token
     * @param eventSourceListener
     * @return
     */
    public void openLLaMA7BStream(OpenLLaMA7BParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, eventSourceListener);
    }

    /**
     * Falcon-7B
     *
     * @param param
     * @param token
     * @return
     */
    public Falcon7BResponse falcon7B(Falcon7BParam param, String token) {
        return (Falcon7BResponse) this.chat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, Falcon7BResponse.class);
    }

    /**
     * Falcon-7B
     *
     * @param param
     * @param token
     * @param eventSourceListener
     * @return
     */
    public void falcon7BStream(Falcon7BParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, eventSourceListener);
    }

    /**
     * Dolly-12B
     *
     * @param param
     * @param token
     * @return
     */
    public Dolly12BResponse dolly12B(Dolly12BParam param, String token) {
        return (Dolly12BResponse) this.chat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, Dolly12BResponse.class);
    }

    /**
     * Dolly-12B
     *
     * @param param
     * @param token
     * @param eventSourceListener
     * @return
     */
    public void dolly12BStream(Dolly12BParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, eventSourceListener);
    }

    /**
     * Llama-2-7b-chat
     * Llama-2-13b-chat
     * Llama-2-70b-chat
     *
     * @param param
     * @param token
     * @param eventSourceListener
     * @return
     */
    public void linlyChineseLLaMAStream(LinlyChineseLLaMAParam param, String token, EventSourceListener eventSourceListener) {
        this.streamChat(String.format(ApiConstant.WITH_API_NAME, param.getApiName()), param, token, eventSourceListener);
    }

    /**
     * visualGLM6B接口
     *
     * @param param
     * @param token
     * @return
     */
    public VisualGLM6BResponse visualGLM6B(VisualGLM6BParam param, String token) {
        String post = HttpUtil.post(String.format(ApiConstant.VISUALGLM_6B, param.getApiName()) + token, JSONUtil.toJsonStr(param));
        return JSONUtil.toBean(post, VisualGLM6BResponse.class);
    }

    /**
     * Stable-Diffusion-XL 文字生成图片
     *
     * @param param
     * @param token 鉴权token
     * @return
     */
    public StableDiffusionXLResponse stableDiffusionXL(StableDiffusionXLParam param, String token) {
        if (param == null) {
            throw new BaiduModelException("参数异常：param不能为空");
        }
        String post = HttpUtil.post(ApiConstant.STABLE_DIFFUSION_XL + token, JSONUtil.toJsonStr(param));
        return JSONUtil.toBean(post, StableDiffusionXLResponse.class);
    }

    /**
     * 发起同步会话
     *
     * @param api   api接口
     * @param param 请求参数
     * @param token 鉴权token
     * @return
     */
    private BaseResponse chat(String api, BaseParam param, String token, Class<? extends BaseResponse> clazz) {
        if (param == null) {
            throw new BaiduModelException("参数异常：param不能为空");
        }
        param.setStream(false);
        String post = HttpUtil.post(api + token, JSONUtil.toJsonStr(param));
        return JSONUtil.toBean(post, clazz);
    }


    /**
     * 发起流式会话
     *
     * @param api   api接口
     * @param param 请求参数
     * @param token 鉴权token
     */
    private void streamChat(String api, BaseParam param, String token, EventSourceListener eventSourceListener) {
        if (Objects.isNull(eventSourceListener)) {
            throw new BaiduModelException("参数异常：EventSourceListener不能为空");
        }
        if (param == null) {
            throw new BaiduModelException("参数异常：param不能为空");
        }
        param.setStream(true);
        try {
            EventSource.Factory factory = EventSources.createFactory(this.okHttpClient);
            ObjectMapper mapper = new ObjectMapper();
            String requestBody = mapper.writeValueAsString(param);
            Request request = new Request.Builder()
                    .url(api + token)
                    .post(RequestBody.create(MediaType.parse(ContentType.JSON.getValue()), requestBody))
                    .build();
            //创建事件
            factory.newEventSource(request, eventSourceListener);
        } catch (JsonProcessingException e) {
            throw new BaiduModelException("请求参数解析是失败！", e);
        }
    }
}
